/*
 * Description:  update Utils in windows
 */
#include "win_update_utils.h"
#include "guest_utils.h"
#include "log.h"
#include <chrono>
#include <ctime>
#include <sstream>
#include <iomanip>
#include "../performance/utils/registry_utils.h"
#include "../performance/utils/perf_utils.h"
#include "../performance/utils/registry_utils_ex.h"

using namespace std::chrono;
using PERF::RegistryUtilsEx;
namespace COMMON {

    bool DoesRegistryKeyExist(HKEY hKeyRoot, const std::string& keyPath)
    {
        HKEY hKey;
        LONG result = RegOpenKeyEx(hKeyRoot, keyPath.c_str(), 0, KEY_READ, &hKey);
        if (result == ERROR_SUCCESS) {
            RegCloseKey(hKey);
            return true;
        } else {
            return false;
        }
    }

    bool CheckRegistryKey(HKEY hKeyRoot, const std::string& keyPath, const std::string& keyName)
    {
        HKEY hKey;
        LONG result = RegOpenKeyEx(hKeyRoot, keyPath.c_str(), 0, KEY_READ, &hKey);
        if (result != ERROR_SUCCESS) {
            LOG_INFO("[win_update_uti]-keyPath not exits. {}", keyPath.c_str());
            return false;
        }
        DWORD valueType;
        BYTE data[1024] = {0};
        DWORD dataSize = sizeof(data);
        result = RegQueryValueEx(hKey, keyName.c_str(), 0, &valueType, data, &dataSize);
        RegCloseKey(hKey);
        return (result == ERROR_SUCCESS);
    }

    bool WinUpdateUtils::CheckWinUpdating()
    {
        //有更新需要重启配置
        std::string keyPath = "SOFTWARE\\Microsoft\\WindowsUpdate\\UX\\StateVariables";
        std::string featureKey = "IsCurrentRebootFeature";
        if (CheckRegistryKey(HKEY_LOCAL_MACHINE, keyPath, featureKey)) {
            LOG_INFO("[win_update_uti]-win update will trigger by reboot.");
            return true;
        }
        LOG_INFO("[win_update_uti]-no win update.");
        return false;
    }
    
    void WinUpdateUtils::DelayUpdate()
    {
        auto now = std::chrono::system_clock::now();
        auto in_time_t = std::chrono::system_clock::to_time_t(now);
        // 将time_t转换为tm结构体，以便格式化
        std::tm* tm_ptr = std::localtime(&in_time_t);
        // 使用stringstream进行格式化输出
        std::stringstream ss1;
        ss1 << std::put_time(tm_ptr, "%Y-%m-%dT%H:%M:%SZ");
        std::string timenow = ss1.str();

        auto one_week_later = now + hours(24 * 7); // 7天，共168小时
        // 将time_point转换为tm结构体，以便格式化为字符串
        std::time_t t = system_clock::to_time_t(one_week_later);
        std::tm* ptm = std::localtime(&t);
        std::stringstream ss2;
        ss2 << std::put_time(ptm, "%Y-%m-%dT%H:%M:%SZ");
        std::string weekLater = ss2.str();
        LOG_INFO("[win_update_uti]-DelayUpdate");
        for (size_t i = 0; i < UPDATES_SETTING_KEYS.size(); i++) {
            auto key = UPDATES_SETTING_KEYS[i];
            LOG_INFO("[win_update_uti]-handle registry key : {}", key);
            std::string value ;
            if (REGISTRY_UTILS.QueryRegSZ(HKEY_LOCAL_MACHINE, BASE_PATH.c_str(), key.c_str(), value)) {
                //备份原有注册表项
                LOG_INFO("[win_update_uti]-backup registry key : {}", key);
                REGISTRY_UTILS.SetRegSZ(HKEY_LOCAL_MACHINE, BAK_PATH.c_str(), key.c_str(), value);
            }
            std::string path = "HKEY_LOCAL_MACHINE\\" + BASE_PATH ;
            if (key.find("Start") != std::wstring::npos) {
                LOG_INFO("[win_update_uti]-set registry key : {}", key);
                std::string cmd = "cmd /q /C \"reg add " + path + " /v " + key + " /t REG_SZ /d " + timenow + " /f\"";
                GuestUtils::ExecCmd(cmd, false);
            } else {
                LOG_INFO("[win_update_uti]-set registry key : {}", key);
                std::string cmd = "cmd /q /C \"reg add " + path + " /v " + key + " /t REG_SZ /d " + weekLater + " /f\"";
                GuestUtils::ExecCmd(cmd, false);
            }
        }
    }

    //还原配置
    void WinUpdateUtils::RestoreSetting()
    {
        LOG_INFO("[win_update_uti]-Restore registry setting");
        for (size_t i = 0; i < UPDATES_SETTING_KEYS.size(); i++) {
            auto key = UPDATES_SETTING_KEYS[i];
            std::string value;
            if (REGISTRY_UTILS.QueryRegSZ(HKEY_LOCAL_MACHINE, BAK_PATH.c_str(), key.c_str(), value)) {
                //还原配置
                LOG_INFO("[win_update_uti]-restore registry key : {}", key);
                REGISTRY_UTILS.SetRegSZ(HKEY_LOCAL_MACHINE, BASE_PATH.c_str(), key.c_str(), value);
            } else {
                LOG_INFO("[win_update_uti]-no need restore registry key : {}", key);
            }
        }
        if (DoesRegistryKeyExist(HKEY_LOCAL_MACHINE, BAK_PATH)) {
            LOG_INFO("[win_update_uti]-delete bak registry key : {}", BAK_PATH);
            std::string path = "HKEY_LOCAL_MACHINE\\" + BAK_PATH;
            std::string cmd = "cmd /q /C \"reg delete " + path + " /f\"";
            GuestUtils::ExecCmd(cmd, false);
        }
    }

    void WinUpdateUtils::UpgradeAgent()
    {
        // 升级执行cmd命令
        UpgradeCMD();
        // 升级执行HKLM
        UpgradeHKLM();
        // 升级执行HKCU
        UpgradeHKCU();
    }

    void WinUpdateUtils::UpgradeHKLM()
    {
        for (RegInfo regInfo : UPGRADE_HKLM_ITEMS) {
            std::string regBakKey = KIDDO_UPGRADE_PATH + regInfo.bakKey;
            std::string value;
            if (REGISTRY_UTILS.QueryRegSZ(HKEY_LOCAL_MACHINE, regBakKey.c_str(), regInfo.regWord.c_str(), value)) {
                LOG_INFO("[win_update_uti]-HKLM has already executed, registry key : {}", regInfo.bakKey);
                continue;
            }
            bool ret;
            if (!regInfo.regValue.empty()) {
                ret = REGISTRY_UTILS.SetRegSZ(HKEY_LOCAL_MACHINE, regInfo.regKey.c_str(),
                regInfo.regWord.c_str(), regInfo.regValue);
            } else {
                ret = REGISTRY_UTILS.SetRegDWORD(HKEY_LOCAL_MACHINE, regInfo.regKey.c_str(),
                regInfo.regWord.c_str(), regInfo.regValueD);
            }
            if (!ret) {
                LOG_ERR("Failed to execute HKLM reg, registry key : {}", regInfo.regKey);
                REGISTRY_UTILS.SetRegSZ(HKEY_LOCAL_MACHINE, regBakKey.c_str(), regInfo.regWord.c_str(), "0");
                continue;
            }
            REGISTRY_UTILS.SetRegSZ(HKEY_LOCAL_MACHINE, regBakKey.c_str(), regInfo.regWord.c_str(), "1");
        }
    }

    void WinUpdateUtils::UpgradeHKCU()
    {
        for (RegInfo regInfo : UPGRADE_HKCU_ITEMS) {
            std::string regBakKey = KIDDO_UPGRADE_PATH + regInfo.bakKey;
            std::string value;
            if (REGISTRY_UTILS.QueryRegSZ(HKEY_LOCAL_MACHINE, regBakKey.c_str(), regInfo.regWord.c_str(), value)) {
                LOG_INFO("[win_update_uti]-HKCU has already executed, registry key : {}", regInfo.bakKey);
                continue;
            }
            HKEY hkcu = HKEY_USERS;
            std::string subKey = PERF_UTILS.GetCurrentUserSid() + "\\" + regInfo.regKey;
            RegistryUtilsEx registryModifier(hkcu, subKey);
            LONG ret;
            if (!regInfo.regValue.empty()) {
                ret = registryModifier.RegSetString(regInfo.regWord, regInfo.regValue);
            } else {
                ret = registryModifier.RegSetDword(regInfo.regWord, regInfo.regValueD);
            }
            if (ret != ERROR_SUCCESS) {
                LOG_ERR("Failed to execute HKCU reg, registry key : {}", regInfo.regKey);
                REGISTRY_UTILS.SetRegSZ(HKEY_LOCAL_MACHINE, regBakKey.c_str(), regInfo.regWord.c_str(), "0");
                continue;
            }
            REGISTRY_UTILS.SetRegSZ(HKEY_LOCAL_MACHINE, regBakKey.c_str(), regInfo.regWord.c_str(), "1");
        }
    }

    void WinUpdateUtils::UpgradeCMD()
    {
        for (RegInfo regInfo : UPGRADE_CMD_ITEMS) {
            std::string regBakKey = KIDDO_UPGRADE_PATH + regInfo.bakKey;
            std::string value;
            if (REGISTRY_UTILS.QueryRegSZ(HKEY_LOCAL_MACHINE, regBakKey.c_str(), regInfo.regWord.c_str(), value)) {
                LOG_INFO("[win_update_uti]-cmd has already executed, registry key : {}", regInfo.bakKey);
                continue;
            }
            std::string cmd = "cmd /q /C " + regInfo.cmdStr;
            bool ret = GuestUtils::ExecCmd(cmd, false);
            if (!ret) {
                LOG_ERR("Failed to execute {}", cmd);
                REGISTRY_UTILS.SetRegSZ(HKEY_LOCAL_MACHINE, regBakKey.c_str(), regInfo.regWord.c_str(), "0");
                continue;
            }
            REGISTRY_UTILS.SetRegSZ(HKEY_LOCAL_MACHINE, regBakKey.c_str(), regInfo.regWord.c_str(), "1");
        }
    }
}
